/*
** runner for wTest
** The MIT License (MIT)
** Copyright (c) <2016> <avaicode>
*/

#ifdef _UNIX
    #define __USE_GNU
    #define _GNU_SOURCE
#endif

#include <stdio.h>
#include <stdlib.h>

#include "src/luaport/luainc.h"

#define MAX_ARG 255

#define SUCCESS 0
#define TIME_LIMIT_EXCEEDED 1
#define MEMORY_LIMIT_EXCEEDED 3
#define RUNTIME_ERROR 4
#define SYSTEM_ERROR 5

#define UNLIMITED -1
#define TIME_INF UNLIMITED
#define MEM_INF UNLIMITED

typedef struct Result {
    int cpuTime;
    unsigned long memory;
    int flag;
    int exitStatus;
} Result;

typedef struct Config {
    int dumpMask;
    int maxCpuTime;
    int maxRealTime;  //Reserved
    unsigned long maxMemory;
    char* path;
    char* inFile;
    char* outFile;
    char* errFile;
    char* args[MAX_ARG];
} Config;

#ifdef _UNIX

#include <errno.h>
#include <sched.h>
#include <unistd.h>
#include <signal.h>
#include <string.h>
#include <sys/time.h>
#include <sys/wait.h>
#include <sys/resource.h>

#define STACK_SIZE (2 * 1024 * 1024)

typedef struct CloneArg {
    Config* configp;
} CloneArg;

int SetTimer(int sec, int ms, int is_cpu_time) {
    struct itimerval timeVal;
    timeVal.it_interval.tv_sec = timeVal.it_interval.tv_usec = 0;
    timeVal.it_value.tv_sec = sec;
    timeVal.it_value.tv_usec = ms * 1000;
    setitimer(is_cpu_time ? ITIMER_VIRTUAL : ITIMER_REAL, &timeVal, NULL);
    return SUCCESS;
}

int ChildProcess(void* cloneArg) {
    Config* config = ((CloneArg*)cloneArg)->configp;
    FILE* inFile = NULL, *outFile = NULL, *errFile = NULL;
    struct rlimit memoryLimit, cpuTimeRlimit;

    if (config->maxMemory != MEM_INF) {
        memoryLimit.rlim_cur = memoryLimit.rlim_max = (rlim_t) (config->maxMemory) * 4 * 1024;
        if (setrlimit(RLIMIT_AS, &memoryLimit) == -1)
            raise(SIGUSR1);
    }

    if (config->maxCpuTime != TIME_INF) {
        SetTimer(config->maxCpuTime / 1000, config->maxCpuTime % 1000, 1);
        cpuTimeRlimit.rlim_cur = cpuTimeRlimit.rlim_max = (config->maxCpuTime + 1000) / 1000;

        if (setrlimit(RLIMIT_CPU, &cpuTimeRlimit) == -1)
            raise(SIGUSR1);
    }

    if (config->inFile != NULL) {
        inFile = fopen(config->inFile, "r");
        if (dup2(fileno(inFile), fileno(stdin)) == -1)
            raise(SIGUSR1);
    }
    if (config->outFile != NULL) {
        outFile = fopen(config->outFile, "w");
        if (dup2(fileno(outFile), fileno(stdout)) == -1)
            raise(SIGUSR1);
    }

    if (config->errFile != NULL) {
        errFile = fopen(config->errFile, "w");
        if (dup2(fileno(errFile), fileno(stderr)) == -1)
            raise(SIGUSR1);
    }

    if (config->dumpMask == 1)
        dup2(fileno(stderr), fileno(stdout));
    if (config->dumpMask == 1 << 1)
        dup2(fileno(stdout), fileno(stderr));

    execvp(config->path, (char* const*)config->args);

#define SAFECLOSE(H) if (H != NULL) fclose(H);
    SAFECLOSE(inFile);
    SAFECLOSE(outFile);
    SAFECLOSE(errFile);

    if (errno == E2BIG)
        return 0;

    raise(SIGUSR1);
    return 1;
}

void Run(Config* config, Result* result) {
    char* stack = NULL;
    int pid, status;
    struct rusage resourceUsage;

    CloneArg cloneArg;
    cloneArg.configp = config;
    stack = malloc(STACK_SIZE);

    pid = clone(ChildProcess, stack + STACK_SIZE, SIGCHLD, (void*)(&cloneArg));
    if (pid < 0) {
        result->flag = SYSTEM_ERROR;
        free(stack);
        return ;
    }

    if (wait4(pid, &status, 0, &resourceUsage) == -1) {
        result->flag = SYSTEM_ERROR;
        free(stack);
        return;
    }

    if (result == NULL) {
        free(stack);
        return ;
    }
    result->flag = SUCCESS;
    result->exitStatus = WEXITSTATUS(status);
    result->cpuTime = (int)(resourceUsage.ru_utime.tv_sec * 1000 +
                            resourceUsage.ru_utime.tv_usec / 1000 +
                            resourceUsage.ru_stime.tv_sec * 1000 +
                            resourceUsage.ru_stime.tv_usec / 1000);
    if (result->cpuTime == 0)
        result->cpuTime = 1;
    result->memory = resourceUsage.ru_maxrss;

    if (WIFSIGNALED(status) != 0) {
        int childSignal = WTERMSIG(status);
        switch (childSignal) {
        case SIGALRM:
            result->flag = TIME_LIMIT_EXCEEDED;
            break;
        case SIGVTALRM:
            result->flag = TIME_LIMIT_EXCEEDED;
            break;
        case SIGSEGV:
            if (result->memory > config->maxMemory)
                result->flag = MEMORY_LIMIT_EXCEEDED;
            else
                result->flag = RUNTIME_ERROR;
            break;
        case SIGUSR1:
            result->flag = SYSTEM_ERROR;
            break;
        default:
            result->flag = RUNTIME_ERROR;
        }
    } else {
        if (result->memory > config->maxMemory)
            result->flag = MEMORY_LIMIT_EXCEEDED;
        if (result->exitStatus != 0)
            result->flag = RUNTIME_ERROR;
    }
    free(stack);
}

#endif

#ifdef _WIN32

#include <windows.h>
#include <psapi.h>

#define MAX_CMD_LEN 1024

typedef struct _THREAD_PARAM {
    HANDLE ProcessHandle;
    int TimeLimit;
} THREAD_PARAM, *LPTHREAD_PARAM;

DWORD WINAPI TimeLimitValidatorThreadProc(LPVOID lpParam);
char* MakeCmdLine(Config* config);

void Run(Config* config, Result* result) {
    HANDLE InputFile = NULL, OutputFile = NULL, ErrFile = NULL;
    int TimeLimit = config->maxCpuTime;
    long MemLimit = config->maxMemory;
    SECURITY_ATTRIBUTES saAttr;
    HANDLE TimeLimitValidator;
    STARTUPINFOA StartupInfo;
    PROCESS_INFORMATION ProcessInfo;
    PROCESS_MEMORY_COUNTERS ProcessMemoryCounters;
    DWORD TimeLimitValidatorThreadID;

    ZeroMemory(&saAttr, sizeof(saAttr));
    saAttr.nLength = sizeof(SECURITY_ATTRIBUTES);
    saAttr.bInheritHandle = TRUE;
    saAttr.lpSecurityDescriptor = NULL;

    ZeroMemory(&ProcessInfo, sizeof(ProcessInfo));
    ZeroMemory(&StartupInfo, sizeof(StartupInfo));
    StartupInfo.cb = sizeof(STARTUPINFO);

    if (config->inFile != NULL)
        InputFile = CreateFile(config->inFile,
                               GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE,
                               &saAttr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (config->outFile != NULL)
        OutputFile = CreateFile(config->outFile,
                                GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE,
                                &saAttr, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
    if (config->errFile != NULL)
        ErrFile = CreateFile(config->errFile,
                             GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE,
                             &saAttr, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
    StartupInfo.hStdInput = InputFile;
    StartupInfo.hStdOutput = OutputFile;
    StartupInfo.hStdError = ErrFile;

    if (config->dumpMask == (1 << 1) + 1) {
        HANDLE t = StartupInfo.hStdOutput;
        StartupInfo.hStdOutput = StartupInfo.hStdError;
        StartupInfo.hStdError = t;
    }
    if (config->dumpMask == 1)
        StartupInfo.hStdOutput = StartupInfo.hStdError;
    if (config->dumpMask == 1 << 1)
        StartupInfo.hStdError = StartupInfo.hStdOutput;

    StartupInfo.dwFlags |= STARTF_USESTDHANDLES;

    char* cmdLine = MakeCmdLine(config);
    CreateProcess(NULL, cmdLine, NULL, NULL,
                  TRUE, CREATE_SUSPENDED, NULL, NULL, (LPSTARTUPINFOA)(&StartupInfo), &ProcessInfo);
    free(cmdLine);

    LPTHREAD_PARAM pData;
    pData = (LPTHREAD_PARAM)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(THREAD_PARAM));
    pData->ProcessHandle = ProcessInfo.hProcess;
    pData->TimeLimit = TimeLimit;
    TimeLimitValidator = CreateThread(NULL, 0, TimeLimitValidatorThreadProc, pData, 0,
                                      &TimeLimitValidatorThreadID);
    ResumeThread(ProcessInfo.hThread);
    WaitForSingleObject(ProcessInfo.hProcess, INFINITE);
    GetProcessMemoryInfo(ProcessInfo.hProcess, &ProcessMemoryCounters, sizeof(ProcessMemoryCounters));
    DWORD WorkingSetUsed = ProcessMemoryCounters.PeakWorkingSetSize / 1024;
    DWORD ThreadExitCode, ExitCode;
    GetExitCodeThread(TimeLimitValidator, &ThreadExitCode);
    CloseHandle(TimeLimitValidator);

    GetExitCodeProcess(pData->ProcessHandle, &ExitCode);
    FILETIME CreateTime, ExitTime, KernelTime, UserTime;
    GetProcessTimes(ProcessInfo.hProcess, &CreateTime, &ExitTime, &KernelTime, &UserTime);
    DWORD TimeUsed = UserTime.dwLowDateTime / 10000;
    if (ThreadExitCode == 2) {
        TimeUsed = TimeLimit + 1;
        ExitCode = -1;
    }

    if (config->dumpMask == 1)
        StartupInfo.hStdOutput = NULL;
    if (config->dumpMask == 1 << 1)
        StartupInfo.hStdError = NULL;
#define SAFECLOSE(H) if (H != NULL) CloseHandle(H);
    SAFECLOSE(StartupInfo.hStdInput);
    SAFECLOSE(StartupInfo.hStdError);
    SAFECLOSE(StartupInfo.hStdOutput);

    if (result != NULL) {
        result->exitStatus = ExitCode;
        result->cpuTime = TimeUsed;
        result->memory = WorkingSetUsed;
        if (TimeUsed > (DWORD)TimeLimit)
            result->flag = TIME_LIMIT_EXCEEDED;
        else if (MemLimit != MEM_INF && (int)WorkingSetUsed > MemLimit)
            result->flag = MEMORY_LIMIT_EXCEEDED;
        else if (ExitCode != 0)
            result->flag = RUNTIME_ERROR;
        else result->flag = SUCCESS;
    }
}

DWORD WINAPI TimeLimitValidatorThreadProc(LPVOID lpParam) {
    LPTHREAD_PARAM pData;
    pData = (LPTHREAD_PARAM)lpParam;
    DWORD SleepTime = pData->TimeLimit - 100 > 0 ?
                      pData->TimeLimit - 100 : pData->TimeLimit;
    if (pData->TimeLimit == TIME_INF)
        SleepTime = 1000;
    while (1) {
        Sleep(SleepTime);
        DWORD ExitCode;
        GetExitCodeProcess(pData->ProcessHandle, &ExitCode);
        if (ExitCode != STILL_ACTIVE)
            return 0;
        FILETIME CreateTime, ExitTime, KernelTime, UserTime, CurrentTime;
        GetSystemTimeAsFileTime(&CurrentTime);
        GetProcessTimes(pData->ProcessHandle, &CreateTime, &ExitTime, &KernelTime, &UserTime);
        DWORD PhysicalTime = (CurrentTime.dwLowDateTime - CreateTime.dwLowDateTime) / 10000;
        DWORD TimeUsed = UserTime.dwLowDateTime / 10000;
        if (pData->TimeLimit != TIME_INF && (int)TimeUsed > pData->TimeLimit) {
            TerminateProcess(pData->ProcessHandle, 0);
            return 1;
        }
        if (pData->TimeLimit != TIME_INF && (int)PhysicalTime > pData->TimeLimit * 2) {
            TerminateProcess(pData->ProcessHandle, 0);
            return 2;
        }
    }
    return 0;
}

char* MakeCmdLine(Config* config) {
    char* ret = (char*)calloc(MAX_CMD_LEN, sizeof(char));
    *ret = '\0';
    char** arg = (char**)config->args;
    for (; *arg != NULL; ++arg) {
        strcat(ret, *arg);
        strcat(ret, " ");
    }
    return ret;
}

#endif

typedef struct ConfItem {
    const char* name;
    int value;
} ConfItem;
static ConfItem conf_item[] = {
    {"_UNLIMITED", UNLIMITED},
    {"_RESULT_SUCCESS", SUCCESS},
    {"_RESULT_TIME_LIMIT_EXCEEDED", TIME_LIMIT_EXCEEDED},
    {"_RESULT_MEMORY_LIMIT_EXCEEDED", MEMORY_LIMIT_EXCEEDED},
    {"_RESULT_RUNTIME_ERROR", RUNTIME_ERROR},
    {"_RESULT_SYSTEM_ERROR", SYSTEM_ERROR}
};
static const int conf_item_size = sizeof(conf_item) / sizeof(ConfItem);

typedef struct ArgMask {
    const char* name;
    int t;  //Should be lua_type
    int nil_accept;
} ArgMask;
static ArgMask arg_mask[] = {
    {"max_cpu_time", LUA_TNUMBER, 0},
    {"max_real_time", LUA_TNUMBER, 1},
    {"max_memory", LUA_TNUMBER, 0},
    {"dump_mask", LUA_TNUMBER, 1},
    {"exe_path", LUA_TSTRING, 0},
    {"input_path", LUA_TSTRING, 1},
    {"output_path", LUA_TSTRING, 1},
    {"error_path", LUA_TSTRING, 1},
    {"args", LUA_TTABLE, 0}
};
static const int arg_mask_size = sizeof(arg_mask) / sizeof(ArgMask);

static int global_Run(lua_State* L);

static const struct luaL_Reg runner[] = {
    {"Run", global_Run},
    {NULL, NULL}
};

#ifdef _UNIX
static char* _strdup(const char* s) {
    char* t = NULL;
    if (s && (t = (char*)malloc(strlen(s) + 1)))
        strcpy(t, s);
    return t;
}
#endif

static void DestoryConfig(Config* config) {
    int i = 0;

    free(config->path);
    free(config->inFile);
    free(config->outFile);
    free(config->errFile);

    for (; i < MAX_ARG; ++i)
        free(config->args[i]);
}

static int lua_checkstringtable(lua_State* L, char** dst, int max_size, int* dst_size) {
    int table_index = lua_gettop(L);
    int table_size = 0;

    lua_pushnil(L);
    while (0 != lua_next(L, table_index)) {
        if (lua_type(L, -1) != LUA_TSTRING)
            return 0;
        if (lua_type(L, -2) != LUA_TNUMBER)
            return 0;

        if (table_size == max_size)
            return 0;

        dst[table_size] = _strdup(lua_tostring(L, -1));
        table_size++;

        lua_pop(L, 1);
    }

    if (dst_size != NULL)
        *dst_size = table_size;

    return 1;
}

static int global_Run(lua_State* L) {
    int i = 0, table_size = 0;
    Config config;
    Result result;
    void* arg_refs[arg_mask_size];

    arg_refs[0] = &config.maxCpuTime;
    arg_refs[1] = &config.maxRealTime;
    arg_refs[2] = &config.maxMemory;
    arg_refs[3] = &config.dumpMask;
    arg_refs[4] = &config.path;
    arg_refs[5] = &config.inFile;
    arg_refs[6] = &config.outFile;
    arg_refs[7] = &config.errFile;
    arg_refs[8] = &config.args;
    memset(&config, 0, sizeof(config));

    luaL_checktype(L, 1, LUA_TTABLE);
    for (; i < arg_mask_size; ++i) {
        lua_pushstring(L, arg_mask[i].name);
        lua_gettable(L, -2);
        if (lua_type(L, -1) != arg_mask[i].t && !(arg_mask[i].nil_accept && lua_type(L, -1) == LUA_TNIL))
            return DestoryConfig(&config), luaL_error(L, "bad argument %s", arg_mask[i].name);


        switch (lua_type(L, -1)) {
        case LUA_TNUMBER:
            *(int*)arg_refs[i] = lua_tointeger(L, -1);
            break;
        case LUA_TSTRING:
            *(char**)arg_refs[i] = _strdup(lua_tostring(L, -1));
            break;
        case LUA_TTABLE:
            if (lua_checkstringtable(L, (char**)arg_refs[i], 256 - 1, &table_size) == 0) {
                DestoryConfig(&config);
                return luaL_error(L, "bad argument %s (item must be a string)", arg_mask[i].name);
            }
            *((char**)arg_refs[i] + table_size) = NULL;
            break;
        case LUA_TNIL:
            if (arg_mask[i].t == LUA_TSTRING)
                *(char**)arg_refs[i] = NULL;
            else if (arg_mask[i].t == LUA_TNUMBER)
                *(int*)arg_refs[i] = 0;
            break;
        }

        lua_pop(L, 1);
    }

    Run(&config, &result);

    lua_newtable(L);
    lua_pushstring(L, "cpu_time"); lua_pushnumber(L, result.cpuTime); lua_rawset(L, -3);
    lua_pushstring(L, "memory"); lua_pushnumber(L, result.memory); lua_rawset(L, -3);
    lua_pushstring(L, "exit_code"); lua_pushnumber(L, result.exitStatus); lua_rawset(L, -3);
    lua_pushstring(L, "flag"); lua_pushnumber(L, result.flag); lua_rawset(L, -3);

    DestoryConfig(&config);
    return 1;
}


int luaopen_runner(lua_State* L) {
    int i = 0;

    luaL_newlib(L, runner);
    for (; i < conf_item_size; ++i) {
        lua_pushstring(L, conf_item[i].name);
        lua_pushnumber(L, conf_item[i].value);
        lua_rawset(L, -3);
    }

    return 1;
}
